import os
from operator import itemgetter
from typing import Union

from langchain_core.output_parsers import StrOutputParser,JsonOutputToolsParser
from langchain_core.runnables import RunnablePassthrough, Runnable, RunnableLambda
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI


class MathTool:
    @tool
    @staticmethod
    def add(x:int,y:int) ->int:
        """
        :summary: 加法函数
        :param x: 第一个数
        :param y: 第二个数
        :return: 两个数的和
        """
        return x+y
    @tool
    @staticmethod
    def sub(x:int,y:int) ->int:
        """
        :summary: 减法函数
        :param x: 第一个数
        :param y: 第二个数
        :return: 两个数的差
        """
        return x-y

    @tool
    @staticmethod
    def mul(x:int,y:int) ->int:
        """
        :summary: 乘法函数
        :param x: 第一个数
        :param y: 第二个数
        :return: 两个数的积
        """
        return x*y

    @tool
    @staticmethod
    def div(x:int,y:int) ->float:
        """
        :summary: 除法函数
        :param x: 第一个数
        :param y: 第二个数
        :return: 两个数的商
        """
        return x/y


def call_tool(tool_invocation: dict) -> Union[str|Runnable]:
    """
    :summary: 调用工具函数
    :param tool_invocation: 工具调用信息
    """
    tool = tool_map[tool_invocation["type"]]
    return RunnablePassthrough.assign(output = itemgetter("args") | tool)

def route(response):
    if len(response["functions"]) > 0:
        return response["functions"]
    else:
        return response["text"]

os.environ["OPENAI_API_KEY"] = os.environ["OPENAI_API_KEY_ZHIHU"]
os.environ["OPENAI_API_BASE"] = os.environ["OPENAI_API_BASE_ZHIHU"]

llm = ChatOpenAI(model="gpt-4o",temperature=0)


tools = [MathTool.add, MathTool.sub, MathTool.mul, MathTool.div]
tool_map = {tool.name: tool for tool in tools}

call_tool_list = RunnableLambda(call_tool).map()

llm_with_tools = (llm.bind_tools(tools) |
    {
       "functions": JsonOutputToolsParser() | call_tool_list,
       "text": StrOutputParser()
    }) | RunnableLambda(route)

result = llm_with_tools.invoke("1024的16倍是多少")
print(result)

result = llm_with_tools.invoke("你是谁")
print(result)